import numpy as np
from typing import List,Dict
from ElementBase import *
from node import Node
from tools import *
# 不支持混合单元
class MeshBlock:
    """网格块Block,用于存储网格块的节点和单元"""
    def __init__(self, name:str, elements:List[elementBase]):
        self.BlockName=name
        self.elements = elements

    def __add__(self,other)->'MeshBlock':
        """test ok"""
        if isinstance(other,MeshBlock):
            if isinstance(self.elements,np.ndarray):
                es1=self.elements.tolist()
            else:
                es1=self.elements
            if isinstance(other.elements,np.ndarray):
                es2=other.elements.tolist()
            else:
                es2=other.elements
            return MeshBlock(name=self.BlockName+'_'+other.BlockName,elements=es1+es2)
        else:
            raise TypeError("unsupported operand type(s) for +: 'MeshBlock' and '{}'".format(type(other).__name__))
    def __radd__(self, other) -> 'MeshBlock':
        """支持 sum 函数,从 0 开始累加"""
        if other == 0:
            return self
        else:
            return self.__add__(other)
    def __str__(self)->str:
        return  f"      MeshBlock:{self.BlockName}\n"+\
                f"      MeshBlock elements num :{self.elemNum}\n"+\
                f"      MeshBlock global dof num :{self.BlockDofNum}\n"
    @property
    def VTK_TYPE(self)->int:
        """vtk单元类型代码"""
        return self.elements[0].__class__.VTK_TYPE # type: ignore
    @property
    def NodesDict(self)->Dict[int,Node]:
        """块的节点对象自典"""
        Nodes:Dict[int,Node]={}
        for Elem in self.elements:
            for nodeOfElem in Elem.NodesList:
                if nodeOfElem.label not in Nodes.keys():
                    Nodes[nodeOfElem.label]=nodeOfElem
        return Nodes
    @property
    def ElemDict(self)->Dict[int,elementBase]:
        """块的单元对象字典"""
        ElemDict:Dict[int,List[elementBase]]={}
        for elem in self.elements:
            if elem.label not in ElemDict.keys():
                ElemDict[elem.label]=elem
        return ElemDict
    @property
    def ElementsType(self)->str:
        return self.elements[0].__class__.__name__ 
    @property
    def DofPerNode(self)->int:
        """块的每个节点的自由度数目"""
        return self.elements[0].NodeDofNum
    @property
    def elemNum(self)->int:    
        return len(self.elements)
    @property
    def BlockDofNum(self)->int:
        """块的全局自由度数目"""
        return self.DofPerNode*self.NodesDict.__len__()
    def Kg(self,globalDofNum:int=0)->np.ndarray:
        """组装block的全局刚度矩阵"""
        kg=np.zeros((globalDofNum,globalDofNum)) 
        for elem in self.elements:
            inds=elem.ElemDofIndexs
            kg[np.ix_(inds,inds)]+=elem.Ke
        return kg
    def Mg(self,globalDofNum:int=0)->np.ndarray:
        """输入全局自由度数目globalDofNum,组装block部分的全局质量矩阵,返回np.ndarray(globalDofNum,globalDofNum)
        
        input:
            globalDofNum:int,全局自由度数目,默认为0,表示使用块的全局自由度数目,即self.BlockDofNum
        output:
            np.ndarray(globalDofNum,globalDofNum),全局质量矩阵
        """
        # 所有单元得到自由度总数
        mg=np.zeros((globalDofNum,globalDofNum)) 
        for elem in self.elements:
            inds=elem.ElemDofIndexs
            mg[np.ix_(inds,inds)]+=elem.Me
        return mg

